/**
 * 
 */
package edu.umd.clip.lm.model.data;

import java.io.*;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.*;

import edu.umd.clip.lm.util.GrowingOutputBuffer;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class OnDiskTrainingDataNodeWriter implements TrainingDataNodeWriter {
	private WritableByteChannel output;
	private static final int bufferSize = 1024*1024;
	private static final int MAX_TEMP_BUFFER = 10*bufferSize;
	private static File tempDir = new File("/tmp");
	
	private int currentData;
	private WritableTrainingDataNode currentNode;
	private Queue<WritableTrainingDataNode> outstandingNodes;

	private enum WritingState {WRITING, WAITING}
	private WritingState currentState;
	private static final boolean DEBUG = false;
	
	private static class TemporaryOutput {
		RandomAccessFile file;
		File filename;
		GrowingOutputBuffer buffer;
		boolean finished;
		
		TemporaryOutput() {
			this.buffer = null;
			this.file = null;
			this.filename = null;
			this.finished = false;
		}
	}
	
	private HashMap<Integer,TemporaryOutput> tempOutput;
	
	/**
	 * @param output
	 */
	public OnDiskTrainingDataNodeWriter(WritableByteChannel output) {
		this.output = new BufferedWritableChannel(output);
		currentNode = null;
		currentData = 0;
		currentState = WritingState.WAITING;
		outstandingNodes = new LinkedList<WritableTrainingDataNode>();
		tempOutput = new HashMap<Integer,TemporaryOutput>();
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.TrainingDataNodeWriter#add(edu.umd.clip.lm.model.training.AbstractTrainingDataNode)
	 */
	@Override
	public void add(WritableTrainingDataNode node) throws IOException {
		outstandingNodes.add(node);
		if (currentState == WritingState.WAITING) {
			selectNextData();
		}
	}
	
	private static int estimateBufferSize(TrainingDataBlock block) {
		int size = block.size() * 100;
		
		// round up to the next multiple of 1024 
		size = (size + 1023) / 1024;
		size *= 1024;
		
		if (size < 1024) 
			return 1024;
		if (size > bufferSize) 
			return bufferSize;
		return size;
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.TrainingDataNodeWriter#write(edu.umd.clip.lm.model.training.AbstractTrainingData, edu.umd.clip.lm.model.training.TrainingDataBlock)
	 */
	@Override
	public void write(WritableTrainingData data, TrainingDataBlock block) throws IOException {
		GrowingOutputBuffer buffer = new GrowingOutputBuffer(estimateBufferSize(block));
		block.write(buffer);
		
		if (currentNode.getData(currentData).getId() == data.getId()) {
			buffer.flip();
			buffer.drainTo(output);
		} else {
			TemporaryOutput temp = tempOutput.get(data.getId());
			if (temp == null) {
				temp = new TemporaryOutput();
				tempOutput.put(data.getId(), temp);
			}
			
			if (temp.buffer == null) {
				temp.buffer = buffer;
			} else {
				temp.buffer.put(buffer);
			}
			
			if (temp.file == null && temp.buffer.size() > MAX_TEMP_BUFFER) {
				createTemporaryFile(temp);
			}
			
			if (temp.file != null) {
				temp.buffer.flip();
				temp.buffer.drainTo(temp.file.getChannel());
				temp.buffer = null;
			}
		}
	}

	private void selectNextData() throws IOException {
		while(true) {
			if (currentState == WritingState.WAITING || currentNode.getNumData() <= currentData + 1) {
				currentNode = outstandingNodes.poll();
				currentData = 0;
				if (currentNode == null) {
					currentState = WritingState.WAITING;
					break;
				}
				currentNode.writeHeader(output);
				currentState = WritingState.WRITING;
			} else {
				++currentData;
			}
			
			if (DEBUG) {
				System.err.printf("NodeWriter: starting node %s\n", currentNode.toString());
			}
			
			TemporaryOutput temp = tempOutput.remove(currentNode.getData(currentData).getId());
			if (temp == null) break;
			
			// copy the temporarily stored data
			if (temp.file != null) {
				long count = temp.file.getChannel().position();
				long written = temp.file.getChannel().transferTo(0, count, output);
				assert(written == count);
				// remove the temp file trying to avoid flushing disk cache
				temp.filename.delete();
				temp.file.setLength(0);
				temp.file.close();
			} else {
				// the data is in the buffer
				temp.buffer.flip();
				temp.buffer.drainTo(output);
			}
			
			if (!temp.finished) break;
		}
	}
	
	private void createTemporaryFile(TemporaryOutput temp) throws IOException {
		try {
			temp.filename = File.createTempFile("on-disk-temp", ".tmp", tempDir);
			temp.filename.deleteOnExit();
			temp.file = new RandomAccessFile(temp.filename, "rw");
		} catch(IOException e) {
			System.err.println("Failed to create temporary file in '"+tempDir+"'");
			throw e;
		}
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.TrainingDataNodeWriter#createData()
	 */
	@Override
	public WritableTrainingData createData() {
		return new WritableTrainingData(this);
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.TrainingDataWriter#finish(edu.umd.clip.lm.model.training.WritableTrainingData)
	 */
	@Override
	public void finish(WritableTrainingData data) {
		if (data == currentNode.getData(currentData)) {
			if (DEBUG) {
				System.err.printf("NodeWriter: finishing node %s\n", currentNode.toString());
			}
			try {
				selectNextData();
			} catch (IOException e) {
				e.printStackTrace();
			}
		} else {
			TemporaryOutput temp = tempOutput.get(data.getId());
			temp.finished = true;
		}
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.TrainingDataWriter#start(edu.umd.clip.lm.model.training.WritableTrainingData)
	 */
	@Override
	public void start(WritableTrainingData data) {
		ByteBuffer buffer = ByteBuffer.allocate(AbstractTrainingData.HEADER_SIZE);
		data.writeHeader(buffer);
		buffer.flip();

		if (data == currentNode.getData(currentData)) {
			try {
				output.write(buffer);
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		} else {
			TemporaryOutput temp = new TemporaryOutput();
			temp.buffer = new GrowingOutputBuffer(bufferSize);
			temp.buffer.put(buffer);
			tempOutput.put(data.getId(), temp);
		}
	}

	public static File getTempDir() {
		return tempDir;
	}

	public static void setTempDir(File tempDir) {
		OnDiskTrainingDataNodeWriter.tempDir = tempDir;
	}

	@Override
	public WritableTrainingDataNode createNode(int nodeId, int numData) {
		WritableTrainingData[] data = new WritableTrainingData[numData];
		for(int i=0; i<numData; ++i) {
			data[i] = createData();
		}
		return new WritableTrainingDataNode(nodeId, data);
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataNodeWriter#sync()
	 */
	@Override
	public void sync() throws IOException {
		// this writer is synchronous already
		return;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataNodeWriter#close()
	 */
	@Override
	public void close() throws IOException {
		// all temporary data should be written down at this point
		output.close();
	}

}
